import tensorflow as tf
import tensorlayer as tl
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

sess = tf.InteractiveSession()

X_train, y_train, X_val, y_val, X_test, y_test = \
                                tl.files.load_mnist_dataset(shape=(-1,784))

             
# 定义模型
x = tf.placeholder(tf.float32, shape=[None, 784], name='x')

network = tl.layers.InputLayer(x, name='input_layer')
network = tl.layers.DropoutLayer(network, keep=0.8, name='drop1')
network = tl.layers.DenseLayer(network, n_units=800,
                                act = tf.nn.relu, name='relu1')
network = tl.layers.DropoutLayer(network, keep=0.5, name='drop2')
network = tl.layers.DenseLayer(network, n_units=800,
                                act = tf.nn.relu, name='relu2')
network = tl.layers.DropoutLayer(network, keep=0.5, name='drop3')
network = tl.layers.DenseLayer(network, n_units=10,
                                act = tf.identity,
                                name='output_layer')

load_params = tl.files.load_npz(path='', name='model.npz')
# sess.run(tf.initialize_all_variables())
tl.files.assign_params(sess, load_params, network)

# print(network)

y = network.outputs
y_op = tf.argmax(tf.nn.softmax(y), 1)


i=0
while i<10000:
    mat_ori=X_test[i].reshape(1,28*28)

    # predict
    print(tl.utils.predict(sess, network, mat_ori, x, y_op))

    # show  
    mat_image=X_test[i].reshape(28,28)*255
    fig = plt.figure()  
    ax = fig.add_subplot(111)  
    ax.imshow(Image.fromarray(mat_image)) 
    plt.show()

    i=i+1
sess.close()